package com.xlh.openai.listener;

import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.xlh.openai.service.ChatRecordsService;
import com.xlh.openai.service.impl.SseEventBuilderImpl;
import com.xlh.openai.dto.ChatCompletionResponseWapperDTO;
import com.xlh.openai.utils.JsonUtils;
import lombok.Setter;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Objects;

/**
 * 描述：OpenAIEventSourceListener
 *
 * @author https:www.unfbx.com
 * @date 2023-02-22
 */
@Slf4j
public class OpenAISSEEventListener extends EventSourceListener {

    private long tokens;

    private SseEmitter sseEmitter;

    @Setter
    private String parentMessageId;

    @Setter
    private String conversationId;

    @Setter
    private String prompt;

    @Setter
    private Integer userId;

    @Setter
    private ChatRecordsService chatRecordsService;


    private String answer = "";

    public OpenAISSEEventListener(SseEmitter sseEmitter) {
        this.sseEmitter = sseEmitter;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void onOpen(EventSource eventSource, Response response) {
        log.info("OpenAI建立sse连接...");
    }

    /**
     * {@inheritDoc}
     */
    @SneakyThrows
    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        tokens += 1;
        if (data.equals("[DONE]")) {
            chatRecordsService.saveConversation( userId, conversationId, prompt, answer, tokens);
            sseEmitter.complete();
            return;
        }
        ObjectMapper mapper = new ObjectMapper();
        // 读取JSON
        ChatCompletionResponse completionResponse = mapper.readValue(data, ChatCompletionResponse.class);
        try {
            if(StrUtil.isNotBlank(completionResponse.getChoices().get(0).getDelta().getContent())){
                answer +=completionResponse.getChoices().get(0).getDelta().getContent();
            }
            ChatCompletionResponseWapperDTO chatCompletionDTO = new ChatCompletionResponseWapperDTO();
            chatCompletionDTO.setId(completionResponse.getId());
            chatCompletionDTO.setParentMessageId(parentMessageId);
            chatCompletionDTO.setText(answer);
            chatCompletionDTO.setConversationId(conversationId);
            chatCompletionDTO.setChatCompletionResponse(completionResponse);
            String responseJson = JsonUtils.toJson(chatCompletionDTO);
            SseEventBuilderImpl sseEventBuilderImpl = new SseEventBuilderImpl();
            sseEventBuilderImpl.data(responseJson, MediaType.APPLICATION_OCTET_STREAM);
            sseEmitter.send(sseEventBuilderImpl);
        } catch (Exception e) {
            log.error("sse信息推送失败！");
            eventSource.cancel();
            e.printStackTrace();
        }
    }


    @Override
    public void onClosed(EventSource eventSource) {
        log.info("流式输出返回值总共{}tokens", tokens() - 2);
        log.info("OpenAI关闭sse连接...");
    }


    @SneakyThrows
    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        if (Objects.isNull(response)) {
            return;
        }
        ResponseBody body = response.body();
        if (Objects.nonNull(body)) {
            log.error("OpenAI  sse连接异常data：{}，异常：{}", body.string(), t);
        } else {
            log.error("OpenAI  sse连接异常data：{}，异常：{}", response, t);
        }
        eventSource.cancel();
    }

    /**
     * tokens
     * @return
     */
    public long tokens() {
        return tokens;
    }
}
